package com.starhub.application.rag.store;

import io.milvus.client.MilvusServiceClient;
import io.milvus.param.R;
import io.milvus.param.dml.*;
import lombok.extern.slf4j.Slf4j;
import io.milvus.grpc.QueryResults;
import io.milvus.grpc.SearchResults;
import io.milvus.grpc.MutationResult;
import dev.langchain4j.data.document.Document;
// LangChain4j imports  
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;

import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.starhub.application.rag.biz.MilvusClientFactory;
import com.starhub.application.rag.constant.RagConstant;
import com.starhub.integration.milvus.langchain.CollectionOperationsExecutor;
import com.starhub.integration.milvus.langchain.Generator;
import com.starhub.integration.milvus.langchain.Mapper;
import com.starhub.utils.Util;

import java.util.*;
import java.util.stream.Collectors;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

/**
 * MilvusStore 类
 * 用于Milvus向量存储的类，提供添加、更新、查询、删除等操作
 * 使用LangChain4j的Metadata对象存储元数据
 * 禁止直接从这里进行向量化执行（受token影响的入口统一在ragservice），这里只用于处理milvus自身的逻辑
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-09
 */
@Slf4j
@Component
public class MilvusStore {
    @Autowired
    private MilvusClientFactory milvusClientFactory;

    // 字段名常量
    private static final String ID_FIELD = RagConstant.LangchainMilvusFields.FIELD_ID;
    private static final String TEXT_FIELD = RagConstant.LangchainMilvusFields.FIELD_TEXT;
    private static final String VECTOR_FIELD = RagConstant.LangchainMilvusFields.FIELD_VECTOR;
    private static final String METADATA_FIELD = RagConstant.LangchainMilvusFields.FIELD_METADATA;

    /**
     * 获取Milvus客户端
     * 
     * @return
     */
    private MilvusServiceClient getClient() {
        MilvusServiceClient milvusClient = milvusClientFactory.getClient(RagConstant.MILVUS_MARK);
        return milvusClient;
    };

    /**
     * 添加TextSegment（LangChain4j风格）
     * 
     * @param content
     * @param embedding
     * @param collectionName
     * @param metaData
     */
    public void add(String content, Embedding embedding, String collectionName, Metadata metaData) {
        Document document = new Document(Util.null2String(content), metaData);
        TextSegment textSegment = document.toTextSegment();
        add(textSegment, embedding, collectionName);
    }

    /**
     * 添加TextSegment（LangChain4j风格）
     * 
     * @param textSegment
     * @param embedding
     * @param collectionName
     * @return
     */
    public String add(TextSegment textSegment, Embedding embedding, String collectionName) {
        List<String> ids = Generator.generateRandomIds(1);
        return add(ids.get(0), textSegment, embedding, collectionName);
    }

    /**
     * 添加TextSegment（指定ID）
     * 
     * @param id
     * @param textSegment
     * @param embedding
     * @param collectionName
     * @return
     */
    public String add(String id, TextSegment textSegment, Embedding embedding, String collectionName) {
        // List<String> ids = Collections.singletonList(id);
        List<String> ids = Generator.generateRandomIds(1);
        List<Embedding> embeddings = Collections.singletonList(embedding);
        List<TextSegment> textSegments = Collections.singletonList(textSegment);
        add(ids, embeddings, textSegments, collectionName);
        return id;
    }

    
    /**
     * 批量添加TextSegment（LangChain4j风格）
     * 
     * @param ids
     * @param embeddings
     * @param textSegments
     * @param collectionName
     */
    public void add(List<String> ids, List<Embedding> embeddings, List<TextSegment> textSegments,
            String collectionName) {
        try {
            List<InsertParam.Field> fields = new ArrayList<>();

            fields.add(new InsertParam.Field(ID_FIELD, ids));
            fields.add(new InsertParam.Field(TEXT_FIELD, Mapper.toScalars(textSegments, ids.size())));
            fields.add(new InsertParam.Field(METADATA_FIELD, Mapper.toMetadataJsons(textSegments, ids.size())));
            fields.add(new InsertParam.Field(VECTOR_FIELD, Mapper.toVectors(embeddings)));

            /*
             * R<MutationResult> result = getClient().insert(
             * InsertParam.newBuilder()
             * .withCollectionName(collectionName)
             * .withFields(fields)
             * .build());
             */

            R<MutationResult> result = CollectionOperationsExecutor.insertWithFlush(getClient(), collectionName,
                    fields);

            if (result.getStatus() != R.Status.Success.getCode()) {
                throw new RuntimeException("Failed to insert data: " + result.getMessage());
            }

        } catch (Exception e) {
            log.error("Failed to add document", e);
            throw new RuntimeException("Failed to add document", e);
        }
    }

    /**
     * 更新指定ID的metadata
     * 
     * @param id
     * @param newMetadata
     * @param collectionName
     * @return
     */
    public boolean updateMetadata(String id, Metadata newMetadata, String collectionName) {
        try {
            // 首先查询现有数据
            DocumentWithMetadata existing = getById(id, collectionName);
            if (existing == null) {
                throw new RuntimeException("Document with id " + id + " not found");
            }

            // 删除原有数据
            R<MutationResult> deleteResult = getClient().delete(
                    DeleteParam.newBuilder()
                            .withCollectionName(collectionName)
                            .withExpr(ID_FIELD + " == \"" + id + "\"")
                            .build());

            if (deleteResult.getStatus() != R.Status.Success.getCode()) {
                throw new RuntimeException("Failed to delete existing data: " + deleteResult.getMessage());
            }

            // 创建新的TextSegment并重新插入
            TextSegment newTextSegment = TextSegment.from(existing.getTextSegment().text(), newMetadata);
            add(id, newTextSegment, existing.getEmbedding(), collectionName);

            return true;
        } catch (Exception e) {
            log.error("Failed to update metadata for id: {}", id, e);
            throw new RuntimeException("Failed to update metadata for id: " + id, e);
        }
    }

    /**
     * 批量更新metadata
     * 
     * @param idToMetadataMap
     * @param collectionName
     * @return
     */
    public boolean updateMetadataBatch(Map<String, Metadata> idToMetadataMap, String collectionName) {
        try {
            for (Map.Entry<String, Metadata> entry : idToMetadataMap.entrySet()) {
                updateMetadata(entry.getKey(), entry.getValue(), collectionName);
            }
            return true;
        } catch (Exception e) {
            log.error("Failed to batch update metadata", e);
            throw new RuntimeException("Failed to batch update metadata", e);
        }
    }

    /**
     * 根据条件更新metadata
     * 
     * @param condition
     * @param newMetadata
     * @param collectionName
     * @return
     */
    public boolean updateMetadataByCondition(String condition, Metadata newMetadata, String collectionName) {
        try {
            List<DocumentWithMetadata> documents = queryByCondition(condition, collectionName);

            Map<String, Metadata> updateMap = new HashMap<>();
            for (DocumentWithMetadata doc : documents) {
                updateMap.put(doc.getId(), newMetadata);
            }

            return updateMetadataBatch(updateMap, collectionName);
        } catch (Exception e) {
            log.error("Failed to update metadata by condition: {}", condition, e);
            throw new RuntimeException("Failed to update metadata by condition: " + condition, e);
        }
    }

    /**
     * 合并metadata（保留现有字段，更新指定字段）
     * 
     * @param id
     * @param additionalMetadata
     * @param collectionName
     * @return
     */
    public boolean mergeMetadata(String id, Metadata additionalMetadata, String collectionName) {
        try {
            DocumentWithMetadata existing = getById(id, collectionName);
            if (existing == null) {
                throw new RuntimeException("Document with id " + id + " not found");
            }

            // 合并metadata
            Metadata mergedMetadata = mergeMetadataObjects(existing.getTextSegment().metadata(), additionalMetadata);

            return updateMetadata(id, mergedMetadata, collectionName);
        } catch (Exception e) {
            log.error("Failed to merge metadata for id: {}", id, e);
            throw new RuntimeException("Failed to merge metadata for id: " + id, e);
        }
    }

    /**
     * 根据metadata字段更新
     * 
     * @param id
     * @param fieldName
     * @param fieldValue
     * @param collectionName
     * @return
     */
    public boolean updateMetadataField(String id, String fieldName, Object fieldValue, String collectionName) {
        try {
            DocumentWithMetadata existing = getById(id, collectionName);
            if (existing == null) {
                throw new RuntimeException("Document with id " + id + " not found");
            }

            // 创建只包含单个字段的metadata
            Map<String, Object> fieldMap = new HashMap<>();
            fieldMap.put(fieldName, fieldValue);
            Metadata fieldMetadata = Metadata.from(fieldMap);

            return mergeMetadata(id, fieldMetadata, collectionName);
        } catch (Exception e) {
            log.error("Failed to update metadata field for id: {}", id, e);
            throw new RuntimeException("Failed to update metadata field for id: " + id, e);
        }
    }

    /**
     * 根据ID查询文档
     * 
     * @param id
     * @param collectionName
     * @return
     */
    public DocumentWithMetadata getById(String id, String collectionName) {
        try {
            R<QueryResults> result = getClient().query(
                    QueryParam.newBuilder()
                            .withCollectionName(collectionName)
                            .withExpr(ID_FIELD + " == \"" + id + "\"")
                            .withOutFields(Arrays.asList(ID_FIELD, TEXT_FIELD, VECTOR_FIELD, METADATA_FIELD))
                            .build());

            if (result.getStatus() != R.Status.Success.getCode()) {
                throw new RuntimeException("Failed to query by id: " + result.getMessage());
            }

            QueryResults queryResults = result.getData();
            if (queryResults.getFieldsDataCount() == 0 ||
                    queryResults.getFieldsData(0).getScalars().getStringData().getDataCount() == 0) {
                return null;
            }

            return parseQueryResult(queryResults, 0);
        } catch (Exception e) {
            log.error("Failed to get document by id: {}", id, e);
            throw new RuntimeException("Failed to get document by id: " + id, e);
        }
    }

    /**
     * 根据条件查询文档
     * 
     * @param condition
     * @param collectionName
     * @return
     */
    public List<DocumentWithMetadata> queryByCondition(String condition, String collectionName) {
        try {
            R<QueryResults> result = getClient().query(
                    QueryParam.newBuilder()
                            .withCollectionName(collectionName)
                            .withExpr(condition)
                            .withOutFields(Arrays.asList(ID_FIELD, TEXT_FIELD, VECTOR_FIELD, METADATA_FIELD))
                            .build());

            if (result.getStatus() != R.Status.Success.getCode()) {
                throw new RuntimeException("Failed to query by condition: " + result.getMessage());
            }

            return parseQueryResults(result.getData());
        } catch (Exception e) {
            log.error("Failed to query by condition: {}", condition, e);
            throw new RuntimeException("Failed to query by condition: " + condition, e);
        }
    }

    /**
     * 向量相似性搜索
     * 
     * @param queryEmbedding
     * @param topK
     * @param filter
     * @param collectionName
     * @return
     */
    public List<SimilarityResult> search(Embedding queryEmbedding, int topK, String filter, String collectionName) {
        return search(queryEmbedding, topK, 0, filter, collectionName);
    }

    /**
     * 向量相似性搜索
     * 
     * @param queryEmbedding
     * @param topK
     * @param filter
     * @param collectionName
     * @return
     */
    public List<SimilarityResult> search(Embedding queryEmbedding, int topK, double minScore, String filter,
            String collectionName) {
        try {
            SearchParam.Builder searchBuilder = SearchParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withMetricType(RagConstant.MILVUS_METRIC_TYPE)
                    .withVectorFieldName(VECTOR_FIELD)
                    .withVectors(Collections.singletonList(queryEmbedding.vectorAsList()))
                    .withTopK(topK)
                    .withParams("{\"nprobe\":10}")
                    .withOutFields(Arrays.asList(ID_FIELD, TEXT_FIELD, METADATA_FIELD));

            if (filter != null && !filter.isEmpty()) {
                searchBuilder.withExpr(filter);
            }

            R<SearchResults> result = getClient().search(searchBuilder.build());

            if (result.getStatus() != R.Status.Success.getCode()) {
                throw new RuntimeException("Failed to search: " + result.getMessage());
            }

            return parseSearchResults(result.getData()).stream()
                    .filter(match -> match.getScore() >= minScore)
                    .collect(Collectors.toList());
        } catch (Exception e) {
            log.error("Failed to perform similarity search", e);
            throw new RuntimeException("Failed to perform similarity search", e);
        }
    }

    /**
     * 删除文档
     * 
     * @param id
     * @param collectionName
     * @return
     */
    public boolean deleteById(String id, String collectionName) {
        try {
            R<MutationResult> result = getClient().delete(
                    DeleteParam.newBuilder()
                            .withCollectionName(collectionName)
                            .withExpr(ID_FIELD + " == \"" + id + "\"")
                            .build());

            return result.getStatus() == R.Status.Success.getCode();
        } catch (Exception e) {
            log.error("Failed to delete document by id: {}", id, e);
            throw new RuntimeException("Failed to delete document by id: " + id, e);
        }
    }

    /**
     * 删除文档
     * 
     * @param condition
     * @param collectionName
     * @return
     */
    public boolean deleteByCondition(String condition, String collectionName) {
        try {
            List<DocumentWithMetadata> list = queryByCondition(condition, collectionName);
            if (list.size() == 0) {
                return true;
            }

            R<MutationResult> result = getClient().delete(
                    DeleteParam.newBuilder()
                            .withCollectionName(collectionName)
                            .withExpr(condition)
                            .build());

            return result.getStatus() == R.Status.Success.getCode();
        } catch (Exception e) {
            log.error("Failed to delete document by condition: {}", condition, e);
            throw new RuntimeException("Failed to delete document by condition: " + condition, e);
        }
    }

    /**
     * 将LangChain4j的Metadata转换为JSON对象
     * 
     * @param metadata
     * @return
     */
    private JsonObject convertMetadataToJson(Metadata metadata) {
        JsonObject json = new JsonObject();
        if (metadata != null) {
            // 获取metadata的所有键值对
            Map<String, Object> metadataMap = metadata.toMap();
            for (Map.Entry<String, Object> entry : metadataMap.entrySet()) {
                if (entry.getValue() instanceof String) {
                    json.addProperty(entry.getKey(), (String) entry.getValue());
                } else if (entry.getValue() instanceof Number) {
                    json.addProperty(entry.getKey(), (Number) entry.getValue());
                } else if (entry.getValue() instanceof Boolean) {
                    json.addProperty(entry.getKey(), (Boolean) entry.getValue());
                } else {
                    json.addProperty(entry.getKey(), entry.getValue().toString());
                }
            }
        }
        return json;
    }

    /**
     * 从JSON字符串构建Metadata对象
     * 
     * @param jsonString
     * @return
     */
    private Metadata convertJsonToMetadata(String jsonString) {
        if (jsonString == null || jsonString.isEmpty()) {
            return new Metadata();
        }

        try {
            JsonObject jsonObject = JsonParser.parseString(jsonString).getAsJsonObject();
            Map<String, Object> metadata = new HashMap<>();
            for (Map.Entry<String, com.google.gson.JsonElement> entry : jsonObject.entrySet()) {
                if (entry.getValue().isJsonPrimitive()) {
                    com.google.gson.JsonPrimitive primitive = entry.getValue().getAsJsonPrimitive();
                    if (primitive.isString()) {
                        metadata.put(entry.getKey(), primitive.getAsString());
                    } else if (primitive.isNumber()) {
                        metadata.put(entry.getKey(), primitive.getAsNumber());
                    } else if (primitive.isBoolean()) {
                        metadata.put(entry.getKey(), primitive.getAsBoolean());
                    }
                } else if (entry.getValue().isJsonNull()) {
                    metadata.put(entry.getKey(), null);
                } else {
                    metadata.put(entry.getKey(), entry.getValue().toString());
                }
            }
            return Metadata.from(metadata);
        } catch (Exception e) {
            e.printStackTrace();
            log.error("Failed to convert json to metadata: {}", jsonString, e);
            return new Metadata();
        }
    }

    /**
     * 合并两个Metadata对象
     * 
     * @param existing
     * @param additional
     * @return
     */
    private Metadata mergeMetadataObjects(Metadata existing, Metadata additional) {
        Map<String, Object> mergedMap = new HashMap<>();

        // 添加现有metadata
        if (existing != null) {
            mergedMap.putAll(existing.asMap());
        }

        // 添加新的metadata（会覆盖相同的key）
        if (additional != null) {
            mergedMap.putAll(additional.asMap());
        }

        return Metadata.from(mergedMap);
    }

    /**
     * 解析查询结果
     * 
     * @param queryResults
     * @return
     */
    private List<DocumentWithMetadata> parseQueryResults(QueryResults queryResults) {
        List<DocumentWithMetadata> documents = new ArrayList<>();

        if (queryResults.getFieldsDataCount() == 0) {
            return documents;
        }

        int rowCount = queryResults.getFieldsData(0).getScalars().getStringData().getDataCount();

        for (int i = 0; i < rowCount; i++) {
            documents.add(parseQueryResult(queryResults, i));
        }

        return documents;
    }

    /**
     * 解析单个查询结果
     * 
     * @param queryResults
     * @param index
     * @return
     */
    private DocumentWithMetadata parseQueryResult(QueryResults queryResults, int index) {
        String id = null;
        String text = null;
        List<Float> embeddingVector = null;
        String metadataJson = null;

        for (io.milvus.grpc.FieldData fieldData : queryResults.getFieldsDataList()) {
            switch (fieldData.getFieldName()) {
                case ID_FIELD:
                    id = fieldData.getScalars().getStringData().getData(index);
                    break;
                case TEXT_FIELD:
                    text = fieldData.getScalars().getStringData().getData(index);
                    break;
                case VECTOR_FIELD:
                    embeddingVector = fieldData.getVectors().getFloatVector().getDataList();
                    break;
                case METADATA_FIELD:
                    metadataJson = fieldData.getScalars().getJsonData().getData(index).toStringUtf8();
                    break;
                default:
                    break;
            }
        }

        if (id == null || text == null || embeddingVector == null || metadataJson == null) {
            log.error("Failed to parse query result: missing required fields");
            throw new RuntimeException("Failed to parse query result: missing required fields");
        }

        Metadata metadata = convertJsonToMetadata(metadataJson);
        TextSegment textSegment = TextSegment.from(text, metadata);
        Embedding embedding = Embedding.from(embeddingVector);

        return new DocumentWithMetadata(id, textSegment, embedding);
    }

    /**
     * 解析搜索结果
     * 
     * @param searchResults
     * @return
     */
    private List<SimilarityResult> parseSearchResults(SearchResults searchResults) {
        List<SimilarityResult> results = new ArrayList<>();

        if (searchResults.getResults().getTopK() == 0) {
            return results;
        }

        for (int i = 0; i < searchResults.getResults().getTopK(); i++) {
            String id = searchResults.getResults().getFieldsData(0).getScalars().getStringData().getData(i);
            String text = searchResults.getResults().getFieldsData(1).getScalars().getStringData().getData(i);
            String metadataJson = searchResults.getResults().getFieldsData(2).getScalars().getJsonData().getData(i)
                    .toStringUtf8();
            float score = searchResults.getResults().getScores(i);

            Metadata metadata = convertJsonToMetadata(metadataJson);
            TextSegment textSegment = TextSegment.from(text, metadata);

            results.add(new SimilarityResult(id, textSegment, score));
        }

        return results;
    }

    /**
     * 数据模型类 - 使用LangChain4j的对象
     */
    public static class DocumentWithMetadata {
        private final String id;
        private final TextSegment textSegment;
        private final Embedding embedding;

        public DocumentWithMetadata(String id, TextSegment textSegment, Embedding embedding) {
            this.id = id;
            this.textSegment = textSegment;
            this.embedding = embedding;
        }

        public String getId() {
            return id;
        }

        public TextSegment getTextSegment() {
            return textSegment;
        }

        public Embedding getEmbedding() {
            return embedding;
        }
    }

    /**
     * 相似性搜索结果模型类
     */
    public static class SimilarityResult {
        private final String id;
        private final TextSegment textSegment;
        private final float score;

        public SimilarityResult(String id, TextSegment textSegment, float score) {
            this.id = id;
            this.textSegment = textSegment;
            this.score = score;
        }

        public String getId() {
            return id;
        }

        public TextSegment getTextSegment() {
            return textSegment;
        }

        public float getScore() {
            return score;
        }
    }
}
